#include "VelocityCrossingMotionSegmenter.h"
#include <MMM/Motion/Plugin/KinematicPlugin/KinematicSensor.h>
#include <MMM/Motion/Plugin/MoCapMarkerPlugin/MoCapMarkerSensor.h>
#include <MMM/Motion/Plugin/ModelPosePlugin/ModelPoseSensor.h>
#include <VirtualRobot/RobotConfig.h>
#include <MMMSimoxTools/MMMSimoxTools.h>

using namespace MMM;

VelocityCrossingMotionSegmenter::VelocityCrossingMotionSegmenter(MotionPtr motion, SegmentationType segType) :
    motion(motion),
    segType(segType)
{
    widget = new VelocityCrossingMotionSegmenterWidget(motion, segType);
}

std::vector<SegmentationPtr> VelocityCrossingMotionSegmenter::segment() {
    return segment(widget->getConfiguration());
}

std::vector<SegmentationPtr> VelocityCrossingMotionSegmenter::segment(VelocityCrossingMotionSegmenterConfigurationPtr configuration) {
    std::vector<SegmentationPtr> segmentation;
    if (configuration) {
        switch (segType) {
        case SegmentationType::MMM:
            {
                KinematicSensorPtr sensor = motion->getSensorByType<KinematicSensor>(KinematicSensor::TYPE);
                std::vector<float> timesteps = sensor->getTimesteps();
                std::vector<std::string> jointNames = sensor->getJointNames();
                std::vector<int> indexes;
                for(std::size_t i=0; i<jointNames.size(); i++){
                    if (configuration->segmentNames.find(jointNames.at(i)) != configuration->segmentNames.end()) {
                        indexes.push_back(i);

                    }
                }

                float minTimestep = 0.0;
                bool lastTimestepSegmented = false;

                for (unsigned int i = 1; i < timesteps.size(); i++) {
                    KinematicSensorMeasurementPtr measurement = sensor->getDerivedMeasurement(timesteps.at(i));
                    KinematicSensorMeasurementPtr measurement2 = sensor->getDerivedMeasurement(timesteps.at(i-1));
                    Eigen::VectorXf jointAngles = measurement->getJointAngles();
                    Eigen::VectorXf jointAngles2 = measurement2->getJointAngles();
                    bool segment = false;
                    int s = 0;
                    for(int indexe : indexes){
                        if(fabs((jointAngles[indexe]-jointAngles2[indexe])/(timesteps.at(i)-timesteps.at(i-1))) <= configuration->velocityThreshold) {
                            s++;
                            if (s >= configuration->segmentThreshold) {
                                if(!lastTimestepSegmented) minTimestep = timesteps.at(i);
                                lastTimestepSegmented = true;
                                segment = true;
                                break;
                            }
                        }
                    }
                    if (!segment && lastTimestepSegmented) {
                        lastTimestepSegmented = false;
                        segmentation.push_back(SegmentationPtr(new Segmentation(minTimestep, timesteps.at(i-1), configuration->keyframeType)));
                    }
                }
                if (lastTimestepSegmented) {
                    segmentation.push_back(SegmentationPtr(new Segmentation(minTimestep, timesteps.at(timesteps.size() - 1), configuration->keyframeType)));
                }
            }
            break;
        case SegmentationType::MMM_MARKER:
            {
                KinematicSensorPtr kSensor = motion->getSensorByType<KinematicSensor>(KinematicSensor::TYPE);
                ModelPoseSensorPtr mpSensor = motion->getSensorByType<ModelPoseSensor>(ModelPoseSensor::TYPE);
                std::vector<std::string> jointNames = kSensor->getJointNames();
                std::vector<float> timesteps = mpSensor->getTimesteps();

                VirtualRobot::RobotPtr robot = MMM::SimoxTools::buildModel(motion->getModel());
                MMM::SimoxTools::updateInertialMatricesFromModels(robot);

                VirtualRobot::RobotPtr robot2 = MMM::SimoxTools::buildModel(motion->getModel());
                MMM::SimoxTools::updateInertialMatricesFromModels(robot2);

                float minTimestep = 0.0;
                bool lastTimestepSegmented = false;
                for (unsigned int i = 1; i < timesteps.size(); i++) {
                    robot->setGlobalPose(mpSensor->getDerivedMeasurement(timesteps.at(i))->getRootPose());
                    Eigen::VectorXf jointAngles = kSensor->getDerivedMeasurement(timesteps.at(i))->getJointAngles();
                    std::map<std::string, float> jointValues;
                    for (unsigned int j = 0; j < jointNames.size(); j++) {
                        jointValues[jointNames.at(j)] = jointAngles(j);
                    }
                    robot->setJointValues(jointValues);

                    robot2->setGlobalPose(mpSensor->getDerivedMeasurement(timesteps.at(i - 1))->getRootPose());
                    Eigen::VectorXf jointAngles2 = kSensor->getDerivedMeasurement(timesteps.at(i - 1))->getJointAngles();
                    std::map<std::string, float> jointValues2;
                    for (unsigned int j = 0; j < jointNames.size(); j++) {
                        jointValues2[jointNames.at(j)] = jointAngles2(j);
                    }
                    robot2->setJointValues(jointValues2);

                    bool segment = false;
                    int s = 0;
                    for (std::string segmentName : configuration->segmentNames) {
                        if(((robot->getSensor(segmentName)->getGlobalPose().col(3).head<3>() - robot2->getSensor(segmentName)->getGlobalPose().col(3).head<3>()).norm())/(timesteps.at(i)-timesteps.at(i-1)) <= configuration->velocityThreshold){
                            s++;
                            if (s >= configuration->segmentThreshold) {
                                if(!lastTimestepSegmented) minTimestep = timesteps.at(i);
                                lastTimestepSegmented = true;
                                segment = true;
                                break;
                            }
                        }
                    }
                    if (!segment && lastTimestepSegmented) {
                        lastTimestepSegmented = false;
                        segmentation.push_back(SegmentationPtr(new Segmentation(minTimestep, timesteps.at(i-1), configuration->keyframeType)));
                    }

                }
                if (lastTimestepSegmented) {
                    segmentation.push_back(SegmentationPtr(new Segmentation(minTimestep, timesteps.at(timesteps.size() - 1), configuration->keyframeType)));
                }
            }
            break;
        case SegmentationType::C3D_MARKER:
            {
                MoCapMarkerSensorPtr sensor = motion->getSensorByType<MoCapMarkerSensor>(MoCapMarkerSensor::TYPE);
                std::vector<float> timesteps = sensor->getTimesteps();

                float minTimestep = 0.0;
                bool lastTimestepSegmented = false;
                for (unsigned int i = 1; i < timesteps.size(); i++) {
                    MoCapMarkerSensorMeasurementPtr measurement = sensor->getDerivedMeasurement(timesteps.at(i));
                    MoCapMarkerSensorMeasurementPtr measurement2 = sensor->getDerivedMeasurement(timesteps.at(i-1));

                    std::map<std::string, Eigen::Vector3f> labeledMarker = measurement->getLabeledMarker();
                    std::map<std::string, Eigen::Vector3f> labeledMarker2 = measurement2->getLabeledMarker();

                    bool segment = false;
                    int s = 0;
                    for (std::string segmentName : configuration->segmentNames) {
                        if(((labeledMarker[segmentName] - labeledMarker2[segmentName]).norm())/(timesteps.at(i)-timesteps.at(i-1)) <= configuration->velocityThreshold){
                            s++;
                            if (s >= configuration->segmentThreshold) {
                                if(!lastTimestepSegmented) minTimestep = timesteps.at(i);
                                lastTimestepSegmented = true;
                                segment = true;
                                break;
                            }
                        }
                    }
                    if (!segment && lastTimestepSegmented) {
                        lastTimestepSegmented = false;
                        segmentation.push_back(SegmentationPtr(new Segmentation(minTimestep, timesteps.at(i-1), configuration->keyframeType)));
                    }

                }
                if (lastTimestepSegmented) {
                    segmentation.push_back(SegmentationPtr(new Segmentation(minTimestep, timesteps.at(timesteps.size() - 1), configuration->keyframeType)));
                }
            }
            break;
        default:
            break;
        }
    }
    return segmentation;
}

void VelocityCrossingMotionSegmenter::setSegmentationType(SegmentationType segType) {
    this->segType = segType;
    widget->setSegmentationType(segType);
}

void VelocityCrossingMotionSegmenter::setMotion(MMM::MotionPtr motion) {
    this->motion = motion;
    widget->setMotion(motion);
}

QWidget* VelocityCrossingMotionSegmenter::getWidget() {
    return widget;
}

std::string VelocityCrossingMotionSegmenter::getName() {
    return NAME;
}
